package BinaryTrees;

import java.util.Comparator;

@SuppressWarnings("unused")
public class AVLTree <E> extends BSTree<E> {
	public AVLTree() {
		this(null);
	}

	public AVLTree (Comparator<E> comparator) {
		super(comparator);
	}
	
	private static class AVLNode <E> extends Node<E> {
		int height = 1;
		
		public AVLNode(E element, Node<E> parent) {
			super(element, parent);
		}
		
		// 更新当前节点高度
		public void updateHeight() {
			int lefHeight = left == null ? 0 : ((AVLNode<E>)left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
			
			height = 1 + Math.max(lefHeight, rightHeight);
		}
		
		// 当前节点是否平衡
		public int balanceFactor() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
			return leftHeight - rightHeight;
		}
		
		// 返回当前节点高度比较高的子节点
		public Node<E> tallerChild() {
			int lefHeight = left == null ? 0 : ((AVLNode<E>)left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
			
			if (lefHeight > rightHeight) {
				return left;
			} else if (lefHeight < rightHeight) {
				return right;
			} else {
				return isLeftChild() ? left : right;
			}
		}
		
		@Override
		public String toString() {
			String parentString = "null";
			if (parent != null) {
				parentString = parent.element.toString();
			}
			return element + "_p(" + parentString + ")_h(" + height + ")";
		}
	}
	
	@Override
	protected void afterAdd(Node<E> node) {
		// TODO Auto-generated method stub
		while ((node = node.parent) != null) {
			if (isBalanced(node)) {
				// 更新高度
				updateHeight(node);
				
			} else {
				// 平衡处理
				rebalance2(node);
				break;
			}
		}
	}
	
	@Override
	protected void afterRemove(Node<E> node) {
		// TODO Auto-generated method stub
		while ((node = node.parent) != null) {
			if (isBalanced(node)) {
				// 更新高度
				updateHeight(node);
			} else {
				rebalance2(node);
			}
		}
	}
	
	@Override
	protected Node<E> createNode(E element, Node<E> parent) {
		// TODO Auto-generated method stub
		return new AVLNode<E>(element, parent);
	}
	
	private boolean isBalanced(Node<E> node) {
		return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
	}
	
	private void updateHeight(Node<E> node) {
		((AVLNode<E>)node).updateHeight();
	}
	
	/**
	 * 恢复平衡方法2
	 * @param grand 高度最低的那个不平衡节点
	 */
	private void rebalance2(Node<E> grand) {
		Node<E> parent = ((AVLNode <E>)grand).tallerChild();
		Node<E> node = ((AVLNode <E>)parent).tallerChild();
		if (parent.isLeftChild()) { // L
			if (node.isLeftChild()) { // LL
				roate(grand, node, node.right, parent, parent.right, grand);
			} else { // LR
				roate(grand, parent, node.left, node, node.right, grand);
			}
			
		} else { // R
			if (node.isLeftChild()) { // RL
				roate(grand, grand, node.left, node, node.right, parent);
			} else { // RR
				roate(grand, grand, parent.left, parent, node.left, node);
			}
		}
	}
	
	// 统一处理四种情况的旋转.
	private void roate(
			Node<E> r,
			Node<E> b, Node<E> c,
			Node<E> d, 
			Node<E> e, Node<E> f) {
		// 让d成为子树根节点
		d.parent = r.parent;
		if (r.isLeftChild()) {
			r.parent.left = d;
		} else if (r.isRightChild()) {
			r.parent.right = d;
		} else {
			root = d;
		}
		
		// 处理b-c
		b.right = c;
		if (c != null) {
			c.parent = b;
		}
		updateHeight(b);
		
		// 处理e-f
		f.left = e;
		if (e != null) {
			e.parent = f;
		}
		
		updateHeight(f);
		
		// 处理b-d-f
		d.left = b;
		d.right = f;
		b.parent = d;
		f.parent = d;
		updateHeight(d);
	}
	
	/**
	 * 恢复平衡方法1
	 * @param grand 高度最低的那个不平衡节点
	 */
	private void rebalance(Node<E> grand) {
		Node<E> parent = ((AVLNode <E>)grand).tallerChild();
		Node<E> node = ((AVLNode <E>)parent).tallerChild();
		
		if (parent.isLeftChild()) { // L
			if (node.isLeftChild()) { // LL
				rotateRight(grand);
			} else { // LR
				rotateLeft(parent);
				rotateRight(grand);
			}
			
		} else { // R
			if (node.isLeftChild()) { // RL
				rotateRight(parent);
				rotateLeft(grand);
			} else { // RR
				rotateLeft(grand);
			}
		}
	}
	
	// 右旋转
	private void rotateRight(Node<E> grand) {
		Node<E> parent = grand.left;
		Node<E> child = parent.right;
		grand.left = child;
		parent.right = grand;
		afterRotate(grand, parent, child);
	}
	
	// 左旋转
	private void rotateLeft(Node<E> grand) {
		Node<E> parent = grand.right;
		Node<E> child = parent.left;
		grand.right = child;
		parent.left = grand;
		afterRotate(grand, parent, child);
	}
	
	private void afterRotate (
			Node<E> grand, 
			Node<E> parent, 
			Node<E> child) {
		// 让parent成为子树的根节点
		parent.parent = grand.parent;
		if (grand.isLeftChild()) {
			grand.parent.left = parent;
		} else if (grand.isRightChild()) {
			grand.parent.right = parent;
		} else {
			// grand是根节点
			root = parent;
		}
		
		// 更新child的parent
		if (child != null) {
			child.parent = grand;
		}
		
		// 更新grand的parent
		grand.parent = parent;
		
		// 更新高度
		updateHeight(grand);
		updateHeight(parent);
	}
}
